package com.gyx.rocsun.common;

import java.util.HashMap;
import java.util.Map;

/**
 * 基数树
 * @author yx.guo
 */
class RadixTree {
    class Node {
        String key;
        Map<Character, Node> children = new HashMap<>();
        boolean isEndOfWord;

        Node(String key) {
            this.key = key;
            this.isEndOfWord = false;
        }
    }

    private Node root;

    public RadixTree() {
        this.root = new Node("");
    }

    public void insert(String word) {
        Node currentNode = root;
        while (!word.isEmpty()) {
            char firstChar = word.charAt(0);
            if (!currentNode.children.containsKey(firstChar)) {
                currentNode.children.put(firstChar, new Node(word));
                currentNode.children.get(firstChar).isEndOfWord = true;
                return;
            }

            Node childNode = currentNode.children.get(firstChar);
            String commonPrefix = getCommonPrefix(childNode.key, word);
            if (commonPrefix.length() < childNode.key.length()) {
                Node newNode = new Node(commonPrefix);
                currentNode.children.put(firstChar, newNode);
                newNode.children.put(childNode.key.charAt(commonPrefix.length()),
                        new Node(childNode.key.substring(commonPrefix.length())));
                newNode.children.get(childNode.key.charAt(commonPrefix.length())).isEndOfWord = childNode.isEndOfWord;
                newNode.children.get(childNode.key.charAt(commonPrefix.length())).children.putAll(childNode.children);
                childNode.key = commonPrefix;
                childNode.isEndOfWord = false;
                currentNode = newNode;
            } else {
                currentNode = childNode;
            }
            word = word.substring(commonPrefix.length());
        }
        currentNode.isEndOfWord = true;
    }

    public boolean search(String word) {
        Node currentNode = root;
        while (!word.isEmpty()) {
            char firstChar = word.charAt(0);
            if (!currentNode.children.containsKey(firstChar)) {
                return false;
            }

            Node childNode = currentNode.children.get(firstChar);
            if (word.startsWith(childNode.key)) {
                word = word.substring(childNode.key.length());
                currentNode = childNode;
            } else {
                return false;
            }
        }
        return currentNode.isEndOfWord;
    }

    private String getCommonPrefix(String str1, String str2) {
        int minLength = Math.min(str1.length(), str2.length());
        for (int i = 0; i < minLength; i++) {
            if (str1.charAt(i) != str2.charAt(i)) {
                return str1.substring(0, i);
            }
        }
        return str1.substring(0, minLength);
    }

    public static void main(String[] args) {
        RadixTree tree = new RadixTree();
        tree.insert("apple");
        tree.insert("app");
        tree.insert("apricot");

        System.out.println(tree.search("apple"));   // true
        System.out.println(tree.search("app"));     // true
        System.out.println(tree.search("apricot")); // true
        System.out.println(tree.search("apex"));    // false
    }
}